# -*- coding: utf-8 -*-
# 根据传入的图片搜索相似图片
# @Time    : 2024/12/9 上午2:14
# @Author  : 赵鑫源
# @FileName: search_image_vectors.py
# @Software: PyCharm
import io
import global_vars  # 全局变量
import json
import my_logging
from PIL import Image

logging = my_logging.LogUtils().get_log()


def process_image(file):
    """
    处理传入的图片文件，返回图片特征向量
    :param file: 图片文件对象
    :return: 图片特征向量
    """
    try:
        # 使用 io.BytesIO 直接处理文件流
        file_stream = io.BytesIO()
        file.save(file_stream)
        file_stream.seek(0)  # 将指针重置到文件开头
        img_path = Image.open(file_stream)
        return global_vars.global_model.forward({'img': img_path})['img_embedding']
    except Exception as e:
        logging.error('图片处理异常：', exc_info=e)
        raise


def format_search_result(sim_i, sim_v):
    """
    格式化搜索结果
    :param sim_i: 相似图片的索引
    :param sim_v: 相似度距离
    :return: 格式化后的结果字典
    """
    uuid = global_vars.index_to_ulid.get(sim_i)  # 通过索引获取 UUID
    if uuid is None:
        return None
    img_info = global_vars.global_i2key.get(uuid)  # 通过 UUID 获取图片信息
    if img_info is None:
        return None

    score = int(sim_v * 1000)
    if score < 2000:
        # 复制并移除 'imgVectors' 键
        img_info_copy = img_info.copy()
        img_info_copy.pop('imgVectors', None)
        return {"score": score, "content": img_info_copy}
    return None


def find_image_vectors(file, k):
    """
    获取传入图片向量，并在内存中查找相似图片
    :param file: 图片文件对象
    :param k: 返回的相似图片数量
    :return: JSON 格式的搜索结果
    """
    try:
        # 获取图片特征向量
        img_embedding = process_image(file)
        # 在索引中搜索相似图片
        distances, indices = global_vars.global_index.search(img_embedding, k)
        result = []

        # 遍历搜索结果，过滤并格式化数据
        for sim_i, sim_v in zip(indices[0], distances[0]):
            formatted_result = format_search_result(sim_i, sim_v)
            if formatted_result:
                result.append(formatted_result)

        return json.dumps({"code": 200, "msg": "TBIR SUCCESS", "data": result}, ensure_ascii=False)

    except Exception as e:
        logging.error('图片检索异常', exc_info=e)
        return json.dumps({"code": 500, "msg": "TBIR FAILED~~~"}, ensure_ascii=False)
